package com.example.sse;

import com.alibaba.fastjson2.JSONObject;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.util.Map;
import java.util.Objects;
import java.util.Timer;
import java.util.TimerTask;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.Consumer;

/**
 * @Author xiaozq
 * @Date 2024/2/23
 * @Description: SSE服务类
 */
@Slf4j
@Component
public class SseEmitterServer{

    @Getter
    private static final ConcurrentHashMap<String, Map<String, SseEmitter>> sseEmitterPool
            = new ConcurrentHashMap<>();

    private static final ConcurrentHashMap<String, Timer>  headerPool
            = new ConcurrentHashMap<>();

    /**
     * 建立连接
     */
    public  SseEmitter connect(String  userCode, String userId){
        log.info("******************开始建立连接*****************");
        //设置超时时间，0表示不过期，默认是30秒，超过时间未完成会抛出异常
        SseEmitter sseemitter = new SseEmitter(0L);
        //注册回调
        sseemitter.onCompletion(completionCallBack(userCode,userId));
        sseemitter.onError(errorCallBack(userCode,userId));
        sseemitter.onTimeout(timeoutCallBack(userCode,userId));
        sseEmitterPool.computeIfAbsent(userCode, k -> new ConcurrentHashMap<>()).put(userId, sseemitter);
        // 开启心跳活跃
        startHeartbeat(sseemitter,userId);
        return sseemitter;
    }

    /**
     * 关闭当前连接
     */
    public void complete(String userCode, String userId){
        Map<String, SseEmitter> map = sseEmitterPool.get(userCode);
        if (map != null)
            map.get(userId).complete();
    }

    /**
     * 关闭所有连接
     */
    public void completeAll(){
        if(!sseEmitterPool.isEmpty()){
            for (Map.Entry<String, Map<String, SseEmitter>> entry : sseEmitterPool.entrySet()) {
                Map<String, SseEmitter> userIdMap = entry.getValue();
                if(!userIdMap.isEmpty()){
                    for (Map.Entry<String, SseEmitter> userIdEntry : userIdMap.entrySet()) {
                        userIdEntry.getValue().complete();
                    }
                }
            }
            sseEmitterPool.clear();
        }
    }

    private  Runnable completionCallBack(String userCode, String userId) {
        return () -> {
            removeUser(userCode,userId);
            log.info("{}结束连接:{}",userCode,userId);
        };
    }

    private  Runnable timeoutCallBack(String userCode, String userId){
        return ()->{
            removeUser(userCode,userId);
            log.error("{}连接超时:{}",userCode,userId);
        };
    }

    private Consumer<Throwable> errorCallBack(String userCode, String userId){
        return throwable -> {
            removeUser(userCode,userId);
            log.error("{}连接异常:{}",userCode,userId);
        };
    }

    /**
     * 推送消息
     */
    public  void sendMessage(String userCode, MessageVO message){
        Map<String, SseEmitter> map = sseEmitterPool.get(userCode);
        if (map != null) {
            for (Map.Entry<String, SseEmitter> entry : map.entrySet()) {
                try {
                    // 发送事件
                    entry.getValue().send(JSONObject.toJSONString(message));
                }catch (Exception e){
                    log.error("{}连接信息:{}, 错误消息:{}",userCode,entry.getKey(),e.getMessage());
                }
            }
        }
    }

    private void removeUser(String userCode, String userId){
        try {
            Map<String, SseEmitter> map = sseEmitterPool.get(userCode);
            if (map != null) {
                map.remove(userId);
                // 如果该用户的所有会话都已关闭，则移除整个映射
                if (map.isEmpty())
                    sseEmitterPool.remove(userCode);
            }
            // 关闭心跳
            stopHeartbeat(userId);
        }catch (Exception e){
            log.error("关闭连接异常{}",e.getMessage());
        }
    }

    /**
     * 开启心跳
     */
    public void startHeartbeat(SseEmitter sseemitter, String userId) {
        Timer heartbeatTimer = new Timer();
        headerPool.put(userId,heartbeatTimer);
        heartbeatTimer.schedule(new TimerTask() {
            @Override
            public void run() {
                if (Objects.nonNull(headerPool.get(userId))) {
                    // 发送心跳:保持长连接
                    try {
                        sseemitter.send("connect active");
                    } catch (Exception e) {
                        log.error("connect active error");
                    }
                }
            }
        }, 25000, 25000);
    }

    /**
     * 关闭心跳
     * @param userId
     */
    public void stopHeartbeat(String userId) {
        Timer timer = headerPool.get(userId);
        if (timer!= null)
            timer.cancel();
        headerPool.remove(userId);
    }
}
